package demo

import org.apache.spark.sql.types.{DataTypes, StructType}
import org.apache.spark.sql.{DataFrame, SparkSession}
import util.SparkUtil

import scala.collection.mutable

/**
 * id,f1,f2,f3,f4,f5
 * 1,40,50,10,2,5
 * 2,80,100,20,4,10
 * 3,121,148,28,6,15
 * 4,35,45,8,2,4
 * 5,70,92,16,4,7.9
 * 6,103,136,23,6.6,12.5
 */
object SimilarityDemo {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkUtil.getSparkSession(this.getClass.getSimpleName)
    val schema: StructType = new StructType()
      .add("id", DataTypes.IntegerType)
      .add("f1", DataTypes.DoubleType)
      .add("f2", DataTypes.DoubleType)
      .add("f3", DataTypes.DoubleType)
      .add("f4", DataTypes.DoubleType)
      .add("f5", DataTypes.DoubleType)
    val frame: DataFrame = spark.read.schema(schema).option("header", value = true).csv("userprofile/data/demo.csv")
    frame.createTempView("df") // 注册视图

    spark.udf.register("eud", eud _) // 将函数注册到sparkSQL中
    spark.udf.register("cos", cos _)
    val frame1: DataFrame = spark.sql(
      """
        |select
        |a.id,
        |b.id,
        |eud(array(a.f1,a.f2,a.f3,a.f4,a.f5), array(b.f1,b.f2,b.f3,b.f4,b.f5)) as eud,
        |cos(array(a.f1,a.f2,a.f3,a.f4,a.f5), array(b.f1,b.f2,b.f3,b.f4,b.f5)) as cos
        |from df a join df b on b.id > a.id
        |""".stripMargin)
    frame1.show(false)
    frame1.printSchema()

  }

  // 计算欧式距离的函数
  def eud(feature1: mutable.WrappedArray[Double], feature2: mutable.WrappedArray[Double]): Double = {
    val tuples: mutable.WrappedArray[(Double, Double)] = feature1.zip(feature2)
    val doubles: mutable.WrappedArray[Double] = tuples.map(tp => Math.pow(tp._2 - tp._1, 2))
    val sum = doubles.sum
    val d2: Double = Math.pow(sum, 0.5)
    1 / (d2 + 1)
  }

  // 计算余弦相似度的函数
  def cos(feature1: mutable.WrappedArray[Double], feature2: mutable.WrappedArray[Double]): Double = {
    val sum1: Double = feature1.map(x => Math.pow(x, 2)).sum
    val sum2: Double = feature2.map(x => Math.pow(x, 2)).sum
    val tuples: mutable.WrappedArray[(Double, Double)] = feature1.zip(feature2)
    val sum: Double = tuples.map(tp => tp._2 * tp._1).sum
    sum / Math.pow(sum1 * sum2, 0.5)
  }
}
